package com.hyz.yuaiagent.config;

import com.hyz.yuaiagent.rag.document.LoveAppDocumentLoader;
import com.hyz.yuaiagent.rag.document.MyKeywordEnricher;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * 04-基于pgsql的向量存储
 * 配置类，用于创建和配置 PgVectorStore 实例
 */
@Configuration
public class PgVectorStoreConfig {

    // 注入 LoveAppDocumentLoader 实例，用于加载文档
    private final LoveAppDocumentLoader loveAppDocumentLoader;
    private final MyKeywordEnricher myKeywordEnricher;

    /**
     * 构造函数注入 LoveAppDocumentLoader
     *
     * @param loveAppDocumentLoader 文档加载器实例
     */
    @Autowired
    public PgVectorStoreConfig(LoveAppDocumentLoader loveAppDocumentLoader, MyKeywordEnricher myKeywordEnricher) {
        this.loveAppDocumentLoader = loveAppDocumentLoader;
        this.myKeywordEnricher = myKeywordEnricher;
    }

    /**
     * 创建并配置 PgVectorStore 实例
     *
     * @param jdbcTemplate            用于操作数据库的 JdbcTemplate 实例
     * @param dashscopeembeddingModel 用于生成向量的 EmbeddingModel 实例
     * @return 配置好的 VectorStore 实例
     */
    @Bean
    public VectorStore pgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeembeddingModel) {
        // 使用 PgVectorStore 构建器配置 PgVectorStore 实例
        PgVectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeembeddingModel)
                .dimensions(1536)                    // 设置向量维度
                .distanceType(COSINE_DISTANCE)       // 设置距离计算类型为余弦距离
                .indexType(HNSW)                     // 设置索引类型为 HNSW
                .initializeSchema(true)              // 自动初始化数据库模式
                .schemaName("public")                // 设置模式名称为 "public"
                .vectorTableName("vector_store")     // 设置向量存储表名为 "vector_store"
                .maxDocumentBatchSize(10000)         // 设置最大文档批处理大小为 10000
                .build();
        List<Document> documents = loveAppDocumentLoader.loaderDocument();
        //AI自动抽取源信息
        List<Document> keywordEnricher = myKeywordEnricher.KeywordEnricher(documents);
/*        //这里不建议一开始初始化的时候就添加 因为这时候 数据库表未初始化 如果第一次运行就在这里添加可能为造成找不到数据库
        vectorStore.add(documents);*/
        vectorStore.add(keywordEnricher);
        return vectorStore;
    }
}
